__author__ = 'moonkey'

import os
import numpy as np
from PIL import Image
from collections import Counter
import pickle as cPickle
import random, math
from data_util.bucketdata import BucketData

class DataGen(object):
    GO = 1
    EOS = 2

    def __init__(self,
                 data_root,
                 annotation_fn,
                 evaluate = False,
                 valid_target_len = float('inf'),
                 img_width_range = (240, 260),
                 word_len = 18):
        """
        args:  
            data_root: 存放dataset的路径，data_base_dir
            annotation_fn: data_path
            evaluate:
            img_width_range: only needed for training set
            word_len: maybe---, length of labels
        returns:
        
        """

        img_height = 32  # 图片h为32
        self.data_root = data_root 
        if os.path.exists(annotation_fn):
            self.annotation_path = annotation_fn  # input_1: dir for labels?
        else:
            self.annotation_path = os.path.join(data_root, annotation_fn)

        if evaluate:  # input_2
            self.bucket_specs = [(int(math.floor(64 / 4)), int(word_len + 2)), (int(math.floor(108 / 4)), int(word_len + 2)),
                                 (int(math.floor(140 / 4)), int(word_len + 2)), (int(math.floor(256 / 4)), int(word_len + 2)),
                                 (int(math.floor(img_width_range[1] / 4)), int(word_len + 2))]
        else:
            # bucket_specs: [(8, 12), (27, 14), (35, 16), (64, 18), (80, 20)]
            # 分母除以4表示CNN中stride了两次?
            self.bucket_specs = [(int(64 / 4), 10 + 2), (int(108 / 4), 12 + 2),
                             (int(140 / 4), 14 + 2), (int(256 / 4), 16 + 2),
                             (int(math.ceil(img_width_range[1] / 4)), word_len + 2)]

        self.bucket_min_width, self.bucket_max_width = img_width_range  # 图片集中的中最小和最大宽度, 200, 320
        self.image_height = img_height  # 32
        self.valid_target_len = valid_target_len

        self.bucket_data = {i: BucketData()
                            for i in range(self.bucket_max_width + 1)}
        self.dictionary = self.query_word()

    def clear(self):
        self.bucket_data = {i: BucketData()
                            for i in range(self.bucket_max_width + 1)}

    def get_size(self):  # 返回数据集length
        with open(self.annotation_path, 'r') as ann_file:
            return len(ann_file.readlines())

    def gen(self, batch_size):
        valid_target_len = self.valid_target_len
        with open(self.annotation_path, 'r') as ann_file:  # images path
            lines = ann_file.readlines()
            random.shuffle(lines)
            for l in lines:  # 每个l为image的路径+labels
                img_path, lex = l.strip().split()  # ["xxxx.jpg", "XXXlabelXXX"]
                try:
                    # img_bw: <class, "np.array">, img with h=32, min < w < max
                    # label: include 'GO' as start, and 'EOS' as end.
                    img_bw, word = self.read_data(img_path, lex)
                    # if valid_target_len < float('inf'):
                    #     word = word[:valid_target_len + 1]  # include GO and EOS
                    width = img_bw.shape[-1]

                    # TODO:resize if > 320
                    b_idx = min(width, self.bucket_max_width)
                    bs = self.bucket_data[b_idx].append(img_bw, word, os.path.join(self.data_root, img_path))  # numbers of images with width=b_idx
                    if bs >= batch_size:
                        b = self.bucket_data[b_idx].flush_out(
                                self.bucket_specs,
                                valid_target_length=valid_target_len,
                                go_shift=1)
                        if b is not None:
                            yield b
                        else:
                            assert False, 'no valid bucket of width %d'%width
                except IOError:
                    pass # ignore error images
        self.clear()

    def read_data(self, img_path, lex):
        '''
        args:
            img_path: image path for reading
            lex: label for the image
        returns:
            img_bw: 3d image
            word: label num
        '''
        assert 0 < len(lex) < self.bucket_specs[-1][1]  # label小于预设的长度,bucket_specs[-1][1]-->word_len
        # L = R * 299/1000 + G * 587/1000 + B * 114/1000
        with open(os.path.join(self.data_root, img_path), 'rb') as img_file:
            img = Image.open(img_file)
            # resize image >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
            w, h = img.size
            aspect_ratio = float(w) / float(h)
            if aspect_ratio < float(self.bucket_min_width) / self.image_height:
                img = img.resize(
                    (self.bucket_min_width, self.image_height),
                    Image.ANTIALIAS)
            elif aspect_ratio > float(
                    self.bucket_max_width) / self.image_height:
                img = img.resize(
                    (self.bucket_max_width, self.image_height),
                    Image.ANTIALIAS)
            elif h != self.image_height:
                img = img.resize(
                    (int(aspect_ratio * self.image_height), self.image_height),
                    Image.ANTIALIAS)
            # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
            
            img_bw = img.convert('L')
            img_bw = np.asarray(img_bw, dtype=np.uint8)
            # img_bw = img_bw[np.newaxis, :]
            img_bw = np.expand_dims(img_bw, axis=0)

        word = [self.GO]
        for c in lex:
            assert c in self.dictionary.keys(), "LABEL [{}] from IMAGE [{}] is not LEGAL.".format(c, img_path)
            word.append(self.dictionary[c])
        word.append(self.EOS)
        word = np.array(word, dtype=np.int32)

        return img_bw, word

    def query_word(self):
        # "0~9", "A B C D E F G H I J K L M N O P Q R S T U V W X Y Z"-"I O S V Z"
        words = ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9",
                 "A", "B", "C", "D", "E", "F", "G", "H", "J", "K",
                 "L", "M", "N", "P", "Q", "R", "T", "U", "W", "X", "Y",
                 "I", "O", "S", "V", "Z"]
        dictionary = {"PADDING": 0, "GO": 1, "EOS": 2}
        for i in range(len(words)):
            dictionary[words[i]] = i + 3
        return dictionary